#include "stdafx.h"
#include "RenderBuffer.h"
#include "BitmapPrinter.h"
#include "VectorMath.h"
#include "Renderer.h"
#include "Scene.h"


using namespace GxLibBasic;
using namespace GxWinGUI;
using namespace RayTracePro;
using namespace GxLibMath;

#define WM_UPDATE_IMAGE WM_USER+101

class MainForm : public Form
{
	Renderer::RenderDevice renderDevice;
	static const int WindowWidth = 1024;
	static const int WindowHeight = 1024;
protected:
	volatile bool completed;
	int GetEncoderClsid(const WCHAR* format, CLSID* pClsid)
	{
		UINT  num = 0;          // number of image encoders
		UINT  size = 0;         // size of the image encoder array in bytes

		ImageCodecInfo* pImageCodecInfo = NULL;

		GetImageEncodersSize(&num, &size);
		if(size == 0)
			return -1;  // Failure

		pImageCodecInfo = (ImageCodecInfo*)(malloc(size));
		if(pImageCodecInfo == NULL)
			return -1;  // Failure

		GetImageEncoders(num, size, pImageCodecInfo);

		for(UINT j = 0; j < num; ++j)
		{
			if( wcscmp(pImageCodecInfo[j].MimeType, format) == 0 )
			{
				*pClsid = pImageCodecInfo[j].Clsid;
				free(pImageCodecInfo);
				return j;  // Success
			}    
		}

		free(pImageCodecInfo);
		return -1;  // Failure
	}
	bool ProcessMessage(WinMessage msg)
	{
		bool rs = Form::ProcessMessage(msg);
		if (msg.message == WM_UPDATE_IMAGE)
		{
			printer.Flush();
			Graphics(this->handle).DrawImage(printer.GetBitmap(), 0, 0);
			if (completed)
			{
				CLSID clsid;
				GetEncoderClsid(L"image/bmp", &clsid);
				printer.GetBitmap()->Save(L"D:\\RenderResult.bmp", &clsid);
				buffer->SavePFM(L"D:\\RenderResult.pfm");
			}
			return true;
		}
		else
		{
			return rs;
		}
	}
private:
	RefPtr<RayTraceSystem> system;
	RefPtr<Scene> scene;
	RefPtr<Renderer> renderer;
	RefPtr<RenderBuffer> buffer;
	BitmapPrinter printer;
	LARGE_INTEGER startTime, lastUpdateTime;
	double invFreq;
	int threads;
	void InitRenderer()
	{
		system = new RayTraceSystem();
		
		buffer = new RenderBuffer();
		buffer->SetSize(WindowWidth, WindowHeight);
		renderer = new Renderer(system.operator->());
		
		renderer->Camera.SetPerspective(WindowWidth, WindowHeight, 50.0f);
		renderer->SetScreenBuffer(buffer.operator->());
		renderer->RenderFinished.Bind(this, &MainForm::RenderCompleted);
		printer.SetResolution(WindowWidth, WindowHeight);
		for (int i = 0; i<WindowWidth; i++)
			for (int j = 0; j<WindowHeight; j++)
				printer.SetPixel(i, j, RayTracePro::Color(0,0,0));
		LARGE_INTEGER freq;
		QueryPerformanceFrequency(&freq);
		invFreq = 1.0 / (double)freq.QuadPart;
		renderDevice = Renderer::RenderDevice::GPU;
		ReportCurrentDevice(renderDevice);
	}

	float viewPhi, viewTheta, viewRadius;
	Vec3 viewCenter;
	void UpdateViewTransform()
	{
		Matrix4 trans, trans2;
	
		Matrix4::RotationY(trans, viewTheta);
		Matrix4::RotationX(trans2, viewPhi);
		trans *= trans2;
		Vec3 dir(0.0f, 0.0f, -1.0f);
		Vec3 transDir;
		trans.Transform(transDir, dir);
		Vec3 pos = viewCenter;
		Vec3 dpos;
		Vec3::Scale(dpos, transDir, -1.0f * viewRadius);
		pos += dpos;
		Matrix4::Translation(trans2, pos.x, pos.y, pos.z);
		trans.LeftMultiply(trans2);
		renderer->Camera.CameraTransform = trans;
	}
	void MainForm_Paint(Object * sender, PaintEventArgs args)
	{
		//args.Graphics->Clear(Gdiplus::Color(0, 0, 0));
		args.Graphics->DrawImage(printer.GetBitmap(), 0, 0);
	}

	void MainForm_MouseDown(Object * sender, MouseEventArgs args)
	{
		renderer->RenderPixel(args.X, args.Y);
	}

	void MainForm_MouseMove(Object * sender, MouseEventArgs args)
	{
	
	}

	void MainForm_MouseUp(Object * sender, MouseEventArgs args)
	{

	}

	void MainForm_KeyPressed(Object * sender, KeyEventArgs args)
	{
		bool render = false;
		if (args.Key == L'0')
		{
			InitScene0();
			render = true;
		}
		else if (args.Key == L'1')
		{
			InitScene1();
			render = true;
		}
		else if (args.Key == L'D' || args.Key == L'd')
		{
			renderDevice = (Renderer::RenderDevice)(((int)renderDevice + 1)%3);
			ReportCurrentDevice(renderDevice);
		}
		else if (args.Key == VK_ESCAPE)
		{
			renderer->Cancel();
		}
		
		if (render)
		{
			printf("Render started...\n");
			renderer->RenderFrameAsync(renderDevice);
		}
	}

	void ReportCurrentDevice(Renderer::RenderDevice dev)
	{
		switch(dev)
		{
		case Renderer::CPU:
			printf("Using CPU only.");
			break;
		case Renderer::GPU:
			printf("Using GPU only.");
			break;
		case Renderer::Both:
			printf("Using both CPU and GPU.");
			break;
		}
		printf("Press [D] to change device.\n");
	}

	void RefreshContent()
	{
		for (int i = 0; i < buffer->Height; i++)
		{
			for (int j = 0; j < buffer->Width; j++)
			{
				RayTracePro::Color c;
				Vec4 & pix = buffer->GetPixel(j, i);
				c.R = min(255, Round(pix.x * 255));
				c.G = min(255, Round(pix.y * 255));
				c.B = min(255, Round(pix.z * 255));
				c.A = 255;
				printer.SetPixel(j, i, c);
			}
		}
		printer.Flush();
		Graphics(this->handle).DrawImage(printer.GetBitmap(), 0, 0);
	}

	void RenderCompleted(bool success, const RenderStatistics & stats)
	{
		if (success)
		{
			printf("Render completed.\n");
			int gpuBlocks = stats.TotalBlocks - stats.BlocksCompletedByCpu;
			printf("Total Blocks: %d (CPU: %d, GPU: %d - %.1f%%)\n", stats.TotalBlocks, stats.BlocksCompletedByCpu, gpuBlocks, gpuBlocks/(float)stats.TotalBlocks*100);
		}
		else
			printf("Render cancelled.\n");
		printf("Time elapsed: %.3fs\n", stats.TimeElapsedInSeconds);
		completed = true;
		RefreshContent();
		PostMessage(this->handle, WM_UPDATE_IMAGE, 0, 0);
	}

public:
	MainForm()
	{
		this->SetMaximizeBox(false);
		this->SetText(L"RayTracePro");
		this->SetBorder(GxWinGUI::fbSingle);
		this->SetClientWidth(WindowWidth);
		this->SetClientHeight(WindowHeight);
		this->OnKeyPressed.Bind(this, &MainForm::MainForm_KeyPressed);
		this->OnPaint.Bind(this, &MainForm::MainForm_Paint);
		this->OnMouseDown.Bind(this, &MainForm::MainForm_MouseDown);
		this->OnMouseMove.Bind(this, &MainForm::MainForm_MouseMove);
		this->OnMouseUp.Bind(this, &MainForm::MainForm_MouseUp);
		InitRenderer();
	}
	~MainForm()
	{
		scene = 0;
		renderer = 0;
		buffer = 0;
	}
private:
	void InitScene0()
	{
		scene = system->CreateScene();
		renderer->SetScene(scene.operator->());
		viewPhi = -0.1f;
		viewTheta = -1.832596f;
		viewRadius = 30.0f;
		//viewCenter = Vec3(-0.0f, 10.0f, 0.0f);
		viewCenter = Vec3(-1280.0f, 650.0f, -95.0f);
		scene->LoadMesh(L"../Demo/Media/sponza.obj");
		scene->meshes.Add(scene->mesh);
		scene->Bvh = new Bvh();
		BvhNode root;
		root.Flags = 1;
		root.Bounds = scene->mesh->Bounds;
		root.Children[0] = 0;
		root.Children[1] = 0;
		root.Parent = BvhNode::Nil;
		GeometryInstance inst;
		Matrix4::CreateIdentityMatrix(inst.Transform);
		Matrix4::CreateIdentityMatrix(inst.TransformInverse);
		scene->Instances.Add(inst);
		scene->Bvh->Nodes.Add(root);
		UpdateViewTransform();
	}
	void InitScene1()
	{
		scene = system->CreateScene();
		renderer->SetScene(scene.operator->());
		viewPhi = 0.0f;//PI/6.0f;
		viewTheta = PI;
		viewRadius = 175.0f;
		viewCenter = Vec3(-0.0f, 22.0f, 60.0f);
		UpdateViewTransform();
		scene->LoadMesh(L"../Demo/Media/bunny.obj");
		scene->meshes.Add(scene->mesh);
		scene->LoadMesh(L"../Demo/Media/box.obj");
		scene->meshes.Add(scene->mesh);
		scene->Bvh = new Bvh();
		scene->Bvh->Nodes.SetCapacity(1024);
		BvhNode rootNode;
		rootNode.Parent = BvhNode::Nil;
		rootNode.Flags = 0;
		rootNode.Bounds.Init();
		scene->Bvh->Nodes.Add(rootNode);
		BvhNode & root = scene->Bvh->Nodes.Begin();
		int randOrder[8] = {3, 6, 1, 7, 0, 4, 5, 2};
		for (int z = 0; z<8; z++)
		{
			scene->Bvh->Nodes.Add(BvhNode());
			BvhNode & zParent = scene->Bvh->Nodes.End();
			zParent.Parent = 0;
			zParent.Flags = 0;
			zParent.Bounds.Init();
			int zParentId = scene->Bvh->Nodes.Count()-1;
			for (int i = 0; i<8; i++)
			{
				scene->Bvh->Nodes.Add(BvhNode());
				BvhNode & parent = scene->Bvh->Nodes.End();
				parent.Bounds.Init();
				int parentId = scene->Bvh->Nodes.Count()-1;
				for (int j = 0; j<8; j++)
				{
					GeometryInstance inst;
					Matrix4 trans;
					Matrix4::Translation(trans, (j-3.5f)*25.0f, (z-3.5f)*22.0f, 20.0f + i * 25.0f);
					Matrix4 rot;
					Matrix4::CreateIdentityMatrix(rot);
					//Matrix4::RotationY(rot, PI/4);
					Matrix4::Multiply(inst.Transform, rot, trans);
					inst.Transform.Inverse(inst.TransformInverse);
					scene->Instances.Add(inst);
					BBox bounds = scene->meshes[i&1]->Bounds;
					inst.Transform.Transform(bounds.Min, Vec3(bounds.Min));
					inst.Transform.Transform(bounds.Max, Vec3(bounds.Max));
					parent.Bounds.Union(bounds);
					parent.Children[j] = scene->Bvh->Nodes.Count();
					BvhNode node;
					node.Flags = 1;
					node.Children[0] = scene->Instances.Count()-1;
					node.Children[1] = i&1;
					node.Parent = parentId;
					node.Bounds = bounds;
					scene->Bvh->Nodes.Add(node);
				}
				parent.Parent = zParentId;
				parent.Flags = 0;
				zParent.Bounds.Union(parent.Bounds);
				zParent.Children[i] = parentId;
			}
			root.Bounds.Union(zParent.Bounds);
			root.Children[z] = zParentId;
		}
		/*
		DeviceList<Ray> rays;
		DeviceList<RayDifferential> rayDiffs;
		DeviceList<DifferentialGeometry> inters;
		rays.SetSize(8);
		rayDiffs.SetSize(8);
		inters.SetSize(8);
		for (int i = 0; i<3; i++)
		{
			rays[i].Direction = Vec3(0.425966f, 0.347872f, 0.835188f);
			rays[i].Origin = Vec3(-0.000015f, 22.000000f, -115.000000f);
			rays[i].ReciprocalDirection.x = 1.0f/rays[i].Direction.x;
			rays[i].ReciprocalDirection.y = 1.0f/rays[i].Direction.y;
			rays[i].ReciprocalDirection.z = 1.0f/rays[i].Direction.z;
			rays[i].tMin = 0.1f;
			rays[i].tMax = FLT_MAX;
		}
		for (int i = 3; i<6; i++)
		{
			rays[i].Direction = Vec3(0.0f, -0.5f, -1.0f);
			Vec3::Normalize(rays[i].Direction, rays[i].Direction);
			rays[i].Origin = Vec3(0.0f, 40.0f, 10.0f);
			rays[i].ReciprocalDirection = Vec3(1.0f/rays[i].Direction.x, 1.0f/rays[i].Direction.y, 1.0f/rays[i].Direction.z);
			rays[i].tMin = 0.1f;
			rays[i].tMax = FLT_MAX;
		}
		for (int i = 6; i<8; i++)
		{
			rays[i].Direction = Vec3(0.0f, -0.2f, -1.0f);
			Vec3::Normalize(rays[i].Direction, rays[i].Direction);
			rays[i].Origin = Vec3(0.0f, 40.0f, 10.0f);
			rays[i].ReciprocalDirection = Vec3(1.0f/rays[i].Direction.x, 1.0f/rays[i].Direction.y, 1.0f/rays[i].Direction.z);
			rays[i].tMin = 0.1f;
			rays[i].tMax = FLT_MAX;
		}
		system->InitRenderThread();
		scene->TraceRay_BVH(inters, rays, rayDiffs);
		*/
	}
};

void TestScan()
{
	DeviceList<int> cpuScan;
	DeviceList<int> toScan, scanOut;
	int arrSize = 256*1024;
	GpuDevice::AllocScanList(cpuScan, arrSize);
	GpuDevice::AllocScanList(toScan, arrSize);
	GpuDevice::AllocScanList(scanOut, arrSize);
	for (int i = 0; i<arrSize; i++)
	{
		toScan[i] = 3;
		cpuScan[i] = toScan[i];
	}
	DebugTimming timer;
	timer.Start("CPU Scan");
	int rs = 0;
	for (int i = 0; i<arrSize; i++)
	{
		int vi = cpuScan[i];
		cpuScan[i] = rs;
		rs += vi;
	}
	timer.End();
	// warm up
	for (int i = 0; i<10; i++)
	{
		GpuDevice::ExclusiveScan(scanOut.DeviceBuffer(false), toScan.DeviceBuffer(true), toScan.Count());
		GpuDevice::GetQueue().finish();
	}
	timer.Start("GPU Scan");
	GpuDevice::ExclusiveScan(scanOut.DeviceBuffer(false), toScan.DeviceBuffer(true), toScan.Count());
	GpuDevice::GetQueue().finish();
	timer.End();
	scanOut.Synchronize();
	
	bool correct = true;
	int i = 0;
	for (i = 0; i<arrSize; i++)
	{
		if (scanOut[i] != cpuScan[i])
		{
			correct = false;
			break;
		}
	}
	if (!correct)
	{
		printf("GPU Scan result wrong at %d, value %d, should be %d, v[i] = %d!\n", i, scanOut[i], cpuScan[i], toScan[i]);
	}
	else
	{
		printf("GPU Scan result correct!\n");
	}
}

int __stdcall _tWinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPTSTR lpCmdLine, int nCmdShow) 
{
	Application::Init(hInstance, hPrevInstance, lpCmdLine, nCmdShow);
	GpuDevice::Initialize();
	//TestScan();
	MainForm * form = new MainForm();
	Application::Run(form);
	GpuDevice::Finalize();
	Application::Dispose();
	return 0;
}

int main()
{
	//AddArrayTest();
	return _tWinMain((HINSTANCE)GetModuleHandle(NULL), NULL, GetCommandLine(), 0);
}


